# copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import paddle
from paddle import nn
from paddle.nn import functional as F
from .det_basic_loss import DiceLoss, BalanceLoss, MaskL1Loss
from .focal_loss import FocalLoss


class FASTLoss(nn.Layer):
    """
    """

    def __init__(self,
                 eps=1e-6,
                 loss_weights=[1.0, 1.0, 1.0],
                 ohem_ratio=3,
                 balance_loss=True,
                 **kwargs):
        super().__init__()
        self.dice_loss = DiceLoss(eps=eps)
        # self.bce_loss = BalanceLoss(
        #     balance_loss=balance_loss,
        #     main_loss_type='BCELoss',
        #     negative_ratio=ohem_ratio)
        self.focal_loss = FocalLoss()

        self.loss_weights = loss_weights

    # def geo_loss(self, pred, target, mask):
    #     eps = 1e-6
    #     pred = paddle.transpose(pred, [0, 2, 3, 1])
    #     target = paddle.transpose(target, [0, 2, 3, 1])
    #     mask = paddle.transpose(mask, [0, 2, 3, 1]) > 0
    #     mask = mask.expand_as(pred)
    #     pred = paddle.masked_select(pred, mask).reshape([-1, 5])
    #     target = paddle.masked_select(target, mask).reshape([-1, 5])
    #     p1, p2, p3, p4, p5 = paddle.split(pred, 5, 1)
    #     t1, t2, t3, t4, t5 = paddle.split(target, 5, 1)

        # area1 = (p1 + p2) * (p3 + p4)
        # area2 = (t1 + t2) * (t3 + t4)
        # w = paddle.maximum(p1, t1) + paddle.minimum(p2, t2)
        # h = paddle.maximum(p3, t3) + paddle.minimum(p4, t4)
        # area_inter = w * h
        # area_union = area1 + area2 - area_inter + eps
        # iou = paddle.clip(area_inter / area_union, min=eps)
        # iou_loss = -paddle.log(iou).mean()
        # assert iou_loss >= 0

        # l1_loss = 0.1 * F.smooth_l1_loss(pred[:, :4], target[:, :4], reduction='mean')

        # angle_loss = 10 * (1. - paddle.cos(p5 - t5)).mean()
        # loss = l1_loss + angle_loss

        # return loss

    def geo_loss(self, pred, target, mask):
        eps = 1e-6
        p1, p2, p3, p4, p5 = paddle.split(pred, 5, 1)
        t1, t2, t3, t4, t5 = paddle.split(target, 5, 1)
        num_pixel = mask.sum() + eps

        area1 = (p1 + p2) * (p3 + p4)
        area2 = (t1 + t2) * (t3 + t4)
        w = paddle.minimum(p1, t1) + paddle.minimum(p2, t2)
        h = paddle.minimum(p3, t3) + paddle.minimum(p4, t4)
        area_inter = w * h
        area_union = area1 + area2 - area_inter + eps
        iou = paddle.clip(area_inter / area_union, min=eps)
        iou_loss = (-paddle.log(iou) * mask).sum() / num_pixel
        assert iou_loss >= 0, iou_loss.numpy()

        l1_loss = F.smooth_l1_loss(pred[:, :4], target[:, :4], reduction='none')
        l1_loss = paddle.sum(l1_loss, 1, keepdim=True)
        l1_loss = 0.01 * (l1_loss * mask).sum() / num_pixel

        angle_loss = 1.0 - paddle.cos(p5 - t5)
        angle_loss = 10 * (angle_loss * mask).sum() / num_pixel
        loss = l1_loss + angle_loss + iou_loss

        return loss

    def forward(self, predicts, labels):
        l_score, l_geo, l_mask = labels[1:]
        f_score = predicts['f_score']
        f_geo = predicts['f_geo']


        # with paddle.no_grad():
        #     m0 = labels[0].numpy()[0]
        #     m0 = denormalize(m0)
        #     m1 = l_score.numpy()[0][0]
        #     m2 = l_geo.numpy()[0]
        #     m3 = l_mask.numpy()[0][0]
        #     import matplotlib.pyplot as plt
        #     fig, axes = plt.subplots(2, 4)
        #     axes = axes.flatten()
        #     axes[0].imshow(m1)
        #     for i in range(1, 6):
        #         axes[i].imshow(m2[i-1])
        #     axes[6].imshow(m3)
        #     axes[7].imshow(m0)
        #     axes[7].imshow(m1, 'hot', alpha=0.5)
        #     plt.show()


        # dice_loss = self.dice_loss(f_score, l_score, l_mask) * self.loss_weights[0]
        # bce_loss = self.bce_loss(f_score, l_score, l_mask) * self.loss_weights[0]

        b_mask = l_mask > 0
        pred_score = paddle.masked_select(f_score, b_mask)
        gt_score = paddle.masked_select(l_score, b_mask)
        dice_loss = self.dice_loss(f_score, l_score, l_mask) * self.loss_weights[0]
        bce_loss = self.focal_loss(pred_score, gt_score) * self.loss_weights[1]

        if l_score.sum() > 0:
            geo_loss = self.geo_loss(f_geo, l_geo, l_score) * self.loss_weights[2]
        else:
            geo_loss = paddle.to_tensor(0.)

        loss = dice_loss + bce_loss + geo_loss
        losses = {
            'loss': loss,
            'dice_loss': dice_loss,
            'bce_loss': bce_loss,
            'geo_loss': geo_loss
        }

        return losses

def denormalize(img):
    import numpy as np
    import cv2

    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])

    img = np.transpose(img, (1, 2, 0))
    img = img[:, :, ::-1]
    img = img * std + mean
    img = np.clip(img, 0, 1).astype('float32')
    img = cv2.resize(img, None, fx=0.5, fy=0.5, interpolation=cv2.INTER_LINEAR)

    return img